from datetime import datetime
import argparse
import os
import numpy as np
import cv2
import torch
from torch import nn, optim
from torch.utils import data as torch_data
from tensorboardX import SummaryWriter
from coolname import generate_slug
from torchvision import transforms

from spair.models import SPAIR
from spair import config as cfg
from spair.dataloader import SimpleScatteredMNISTDataset
from spair import debug_tools
from spair import metric


dt = datetime.today().strftime('%b-%d') + '-' + generate_slug(2)
run_log_path = 'logs_v2/%s' % dt
writer = SummaryWriter(run_log_path)
print('log path:', run_log_path)

parser = argparse.ArgumentParser()
parser.add_argument('--gpu', help='Enable GPU use', action='store_true',default=1)
args = parser.parse_args()
if args.gpu:
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
    DEVICE = torch.device("cpu")


def train():
    if not os.path.exists(os.path.join('./','imgs/')):
        os.mkdir(os.path.join('./', 'imgs/'))
    image_shape = cfg.INPUT_IMAGE_SHAPE
    torch.manual_seed(3)
    spair_net = SPAIR(image_shape, writer, DEVICE).to(DEVICE)
    params = spair_net.parameters()
    spair_optim = optim.Adam(params, lr=1e-4)
    for epoch in range(200000):
        x_image = cv2.imread('concat_1.png')
        x_image = transforms.ToTensor()(x_image)
        x_image = x_image.unsqueeze(0)
        x_image = torch.cat((x_image, x_image),dim=0)
        x_image = x_image.to(DEVICE)#the image
        spair_optim.zero_grad()
        loss, out_img, z_where, z_pres = spair_net(x_image, epoch)
        loss.backward(retain_graph = True)
        spair_optim.step()
        # logging stuff
        image_out = out_img[0]
        image_in = x_image[0]
        combined_image = torch.cat([image_in, image_out], dim=2)
        combined_image = combined_image.permute(1, 2, 0)
        combined_image = (combined_image.cpu().detach().numpy() * 255.0).astype(np.uint8)
        print('epoch',epoch,'is done')
        if epoch % 100 == 0:
            path = 'imgs/output'+str(epoch)+'.jpg'
            cv2.imwrite(path,combined_image)
            torch.save(spair_net.state_dict(), 'checkpoints/checkpoint'+str(epoch))
    torch.cuda.empty_cache()


if __name__ == '__main__':
    train()